# -*- coding: UTF-8 -*-
import socket
import threading
import logging

#存储已注册的 Worker 节点及其状态，False 表示空闲，True 表示忙碌
registered_workers = []


def handle_worker_registration(worker_socket, worker_info):
    try:
        # 初始状态为空闲
        registered_workers.append((worker_socket, False))
        print(f"Worker {worker_info} 已注册")
        worker_socket.sendall(b'OK')
    except Exception as e:
        print(f"Worker 注册出错: {e}")


def handle_client(client_socket):
    try:
        available_workers = [worker for worker, busy in registered_workers if not busy]
        if not available_workers:
            client_socket.sendall(b'No available workers')
            client_socket.close()
            return

        # 接收图片数据长度
        data_length = int(client_socket.recv(1024).decode())
        print(data_length)
        # 发送确认信息
        client_socket.sendall(b'OK')

        # 接收实际图片数据
        image_data = b''
        while len(image_data) < data_length:
            chunk = client_socket.recv(4096)
            if not chunk:
                break
            image_data += chunk

        # 选择一个空闲的 Worker 节点来处理任务
        worker_socket = available_workers[0]
        for i, (ws, busy) in enumerate(registered_workers):
            if ws == worker_socket:
                # 标记该 Worker 为忙碌
                registered_workers[i] = (ws, True)
                break

        try:
            # 发送图片数据长度
            worker_socket.sendall(str(len(image_data)).encode())
            # 等待 Worker 确认接收长度信息
            worker_socket.recv(1024)
            # 发送实际图片数据
            worker_socket.sendall(image_data)

            # 接收处理结果
            res = worker_socket.recv(1024)

            # 发送处理结果给 Client
            client_socket.sendall(res)

        except Exception as e:
            print(f"Worker 处理错误: {e}")
            logging.exception(e)
        finally:
            for i, (ws, busy) in enumerate(registered_workers):
                if ws == worker_socket:
                    # 标记该 Worker 为空闲
                    registered_workers[i] = (ws, False)
                    break

    except Exception as e:
        print(f"处理 Client 连接时出错: {e}")
        logging.exception(e)
    finally:
        client_socket.close()


def start_master(master_port):
    master_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    master_socket.bind(('0.0.0.0', master_port))
    master_socket.listen(5)

    print(f"Master 节点正在监听 0.0.0.0:{master_port}")

    while True:
        client_or_worker_socket, address = master_socket.accept()
        first_data = client_or_worker_socket.recv(1024).decode()
        print("第一次接受的数据： ", first_data)
        if first_data.startswith('REGISTER'):
            # 处理 Worker 注册
            handle_worker_registration(client_or_worker_socket, first_data)
        else:
            # 处理 Client 请求
            print('Client come in.')
            client_thread = threading.Thread(target=handle_client, args=(client_or_worker_socket,))
            client_thread.start()


# 使用PYTHONUNBUFFERED=1环境变量解决，docker运行容器python的print不输出
if __name__ == "__main__":
    master_port = 8888
    start_master(master_port)
